#include "esolver_fp.h"

#include "source_estate/cal_ux.h"
#include "source_estate/module_charge/symmetry_rho.h"
#include "source_estate/read_pseudo.h"
#include "source_hamilt/module_ewald/H_Ewald_pw.h"
#include "source_hamilt/module_vdw/vdw.h"
#include "source_io/cif_io.h"
#include "source_io/cube_io.h" // use write_vdata_palgrid
#include "source_io/json_output/init_info.h"
#include "source_io/json_output/output_info.h"
#include "source_io/output_log.h"
#include "source_io/print_info.h"
#include "source_io/rhog_io.h"
#include "source_io/module_parameter/parameter.h"

#include "source_pw/module_pwdft/setup_pwrho.h" // mohan 20251005
#include "source_hamilt/module_xc/xc_functional.h" // mohan 20251005
#include "source_io/ctrl_output_fp.h"
#include "source_io/write_init.h" // write_chg_init, write_pot_init

namespace ModuleESolver
{

ESolver_FP::ESolver_FP()
{
}

ESolver_FP::~ESolver_FP()
{
	//****************************************************
	// do not add any codes in this deconstructor funcion
	//****************************************************
    // mohan add 20251005
    pw::teardown_pwrho(this->pw_rho_flag, PARAM.globalv.double_grid, this->pw_rho, this->pw_rhod);

	delete this->pelec;
}

void ESolver_FP::before_all_runners(UnitCell& ucell, const Input_para& inp)
{
    ModuleBase::TITLE("ESolver_FP", "before_all_runners");

    //! read pseudopotentials
    elecstate::read_pseudo(GlobalV::ofs_running, ucell);

    // setup pw_rho, pw_rhod, pw_big, sf, and read_pseudopotentials
    pw::setup_pwrho(ucell, PARAM.globalv.double_grid, this->pw_rho_flag, 
      this->pw_rho, this->pw_rhod, this->pw_big, 
      this->classname, inp);

    // setup structure factors
    this->sf.set(this->pw_rhod, inp.nbspline);

    // write geometry file
    ModuleIO::CifParser::write(PARAM.globalv.global_out_dir + "STRU.cif",
      ucell, "# Generated by ABACUS ModuleIO::CifParser", "data_?");

    // init charge extrapolation
    this->CE.Init_CE(inp.nspin, ucell.nat, this->pw_rhod->nrxx, inp.chg_extrap);

    return;
}

void ESolver_FP::after_scf(UnitCell& ucell, const int istep, const bool conv_esolver)
{
    ModuleBase::TITLE("ESolver_FP", "after_scf");

    //! Output convergence information
    ModuleIO::output_convergence_after_scf(conv_esolver, this->pelec->f_en.etot);

    //! Write Fermi energy
    ModuleIO::output_efermi(conv_esolver, this->pelec->eferm.ef);

    //! Update delta_rho for charge extrapolation
    CE.update_delta_rho(ucell, &(this->chr), &(this->sf));

    //! print out charge density, potential, elf, etc.
	ModuleIO::ctrl_output_fp(ucell, this->pelec, this->pw_big, this->pw_rhod, 
			this->chr, this->solvent, this->Pgrid, istep); 

}

void ESolver_FP::before_scf(UnitCell& ucell, const int istep)
{
    ModuleBase::TITLE("ESolver_FP", "before_scf");

    // if the cell has changed
    if (ucell.cell_parameter_updated)
    {
        // only G-vector and K-vector are changed due to the change of lattice
        // vector FFT grids do not change!!
        this->pw_rho->initgrids(ucell.lat0, ucell.latvec, pw_rho->nx, pw_rho->ny, pw_rho->nz);
        this->pw_rho->collect_local_pw();
        this->pw_rho->collect_uniqgg();

        // if double grid used in USPP, update related quantities in dense grid
        if (PARAM.globalv.double_grid)
        {
            this->pw_rhod->initgrids(ucell.lat0, ucell.latvec, pw_rhod->nx, pw_rhod->ny, pw_rhod->nz);
            this->pw_rhod->collect_local_pw();
            this->pw_rhod->collect_uniqgg();
        }

        // reset local pseudopotentials
        this->locpp.init_vloc(ucell, this->pw_rhod);
        ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "LOCAL POTENTIAL");

        // perform symmetry analysis
        if (ModuleSymmetry::Symmetry::symm_flag == 1)
        {
            ucell.symm.analy_sys(ucell.lat, ucell.st, ucell.atoms, GlobalV::ofs_running);
            ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "SYMMETRY");
        }

        // reset k-points
        KVectorUtils::set_after_vc(kv, PARAM.inp.nspin, ucell.G);
        ModuleBase::GlobalFunc::DONE(GlobalV::ofs_running, "INIT K-POINTS");
    }

    // charge extrapolation
    if (ucell.ionic_position_updated)
    {
        this->CE.update_all_dis(ucell);
        this->CE.extrapolate_charge(&this->Pgrid, ucell, &this->chr, &this->sf,
                                    GlobalV::ofs_running, GlobalV::ofs_warning);
    }

    //! calculate D2 or D3 vdW
    auto vdw_solver = vdw::make_vdw(ucell, PARAM.inp, &(GlobalV::ofs_running));
    if (vdw_solver != nullptr)
    {
        this->pelec->f_en.evdw = vdw_solver->get_energy();
    }

    //! calculate ewald energy
    if (!PARAM.inp.test_skip_ewald)
    {
        this->pelec->f_en.ewald_energy = H_Ewald_pw::compute_ewald(ucell, this->pw_rhod, this->sf.strucFac);
    }

    //! set direction of magnetism, used in non-collinear case 
    elecstate::cal_ux(ucell);

    //! output the initial charge density and potential
    ModuleIO::write_chg_init(ucell, this->Pgrid, this->chr, this->pelec->eferm, istep, PARAM.inp);
//    ModuleIO::write_pot_init(ucell, this->Pgrid, this->pelec, istep, PARAM.inp); 

    return;
}

void ESolver_FP::iter_finish(UnitCell& ucell, const int istep, int& iter, bool& conv_esolver)
{
    //! output charge density in G-space, or if available, kinetic energy density in G-space
    if (PARAM.inp.out_chg[0] != -1)
    {
        if (iter % PARAM.inp.out_freq_elec == 0 || iter == PARAM.inp.scf_nmax || conv_esolver)
        {
            for (int is = 0; is < PARAM.inp.nspin; is++)
            {
                this->pw_rhod->real2recip(this->chr.rho_save[is], this->chr.rhog_save[is]);
            }
            ModuleIO::write_rhog(PARAM.globalv.global_out_dir + PARAM.inp.suffix + "-CHARGE-DENSITY.restart",
                                 PARAM.globalv.gamma_only_pw || PARAM.globalv.gamma_only_local,
                                 this->pw_rhod,
                                 PARAM.inp.nspin,
                                 ucell.GT,
                                 this->chr.rhog_save,
                                 GlobalV::MY_POOL,
                                 GlobalV::RANK_IN_POOL,
                                 GlobalV::NPROC_IN_POOL);

            if (XC_Functional::get_ked_flag())
            {
                std::vector<std::complex<double>> kin_g_space(PARAM.inp.nspin * this->chr.ngmc, {0.0, 0.0});
                std::vector<std::complex<double>*> kin_g;
                for (int is = 0; is < PARAM.inp.nspin; is++)
                {
                    kin_g.push_back(kin_g_space.data() + is * this->chr.ngmc);
                    this->pw_rhod->real2recip(this->chr.kin_r_save[is], kin_g[is]);
                }
                ModuleIO::write_rhog(PARAM.globalv.global_out_dir + PARAM.inp.suffix + "-TAU-DENSITY.restart",
                                     PARAM.globalv.gamma_only_pw || PARAM.globalv.gamma_only_local,
                                     this->pw_rhod,
                                     PARAM.inp.nspin,
                                     ucell.GT,
                                     kin_g.data(),
                                     GlobalV::MY_POOL,
                                     GlobalV::RANK_IN_POOL,
                                     GlobalV::NPROC_IN_POOL);
            }
        }
    }
}

void ESolver_FP::after_all_runners(UnitCell& ucell)
{
    // print out the final total energy
    GlobalV::ofs_running << "\n --------------------------------------------" << std::endl;
    GlobalV::ofs_running << std::setprecision(16);
    GlobalV::ofs_running << " !FINAL_ETOT_IS " << this->pelec->f_en.etot * ModuleBase::Ry_to_eV << " eV" << std::endl;
    GlobalV::ofs_running << " --------------------------------------------\n\n" << std::endl;

}

} // namespace ModuleESolver
